package counter.com.myservice;

import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import counter.com.baomidou.controller.LoudongbuginfoController;
import counter.com.baomidou.controller.LoudongscaninfoController;
import counter.com.baomidou.entity.*;
import counter.com.baomidou.service.*;
import counter.com.common.MyTool;
import counter.com.common.ResultWebSocket;
import counter.com.common.StoreManager;
import counter.com.configuration.LoudongConfiguration;
import lombok.Data;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.Serializable;
import java.time.LocalDateTime;
import java.util.Collections;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;

@Component
public class PollingService {
    public static final int TIMESPAN = 2000;
    public static final int FAST_INCREASE_PROGRESS = 500;
    public static final int FULLPROGRESS = 10000;

    public static final int STATUS_STOP = 0;
    public static final int STATUS_GOING = 1;
    public static final int STATUS_FINISH = 2;


    List<Loudongscaninfo> scaninfoList;
    List<Attackinfo> attackinfoList;

    @Autowired
    ILoudongtargetinfoService targetService;

    @Autowired
    ILoudongscaninfoService scanInfoService;

    @Autowired
    IAttackinfoService attackinfoService;

    @Autowired
    IMethodinfoService methodinfoService;

    @Autowired
    IActivityinfoService activityinfoService;

    @Autowired
    ILoudongbuginfoService bugService;

    @Autowired
    IPresetbuginfoService presetbugService;

    @Autowired
    LoudongConfiguration configuration;

    private static BugGenerator instance;

    private BugGenerator getBugGenerator() {
        if (instance == null) {
            instance = new BugGenerator();
        }
        return instance;
    }

    public static String[] bugNames = {"nginx SPDY heap buffer overflow", "HTML form whihout CSRF protection", "OPTIONS method is enabled", "Broken links"};

    /**
     * 轮询服务的构造函数
     */
    @Autowired
    public PollingService() {
        System.out.println("轮询服务启动，时间间隔是" + TIMESPAN + "ms ......");
        scaninfoList = new LinkedList<>();
        attackinfoList = new LinkedList<>();
        System.out.println("从数据库中载入未完成的扫描数据....");
//        initDataFromServer();
        System.out.println("....载入完成");
        Runnable runnable = new Runnable() {
            @Override
            public void run() {
                while (true) {
                    try {
                        if (!getAttackInfoList().isEmpty()) {
                            processOfAttackInfoList();
                        }
                        if (!getScanInfoList().isEmpty()) {
                            processOfScanInfoList();
                        }
                        Thread.sleep(TIMESPAN);
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                }
            }
        };
        Thread thread = new Thread(runnable);
        thread.start();
    }

    private void initDataFromServer() {
        LambdaQueryWrapper<Loudongscaninfo> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(Loudongscaninfo::getStatus, STATUS_GOING);
        System.out.println(scanInfoService);
        List<Loudongscaninfo> list = scanInfoService.list(wrapper);
        scaninfoList.addAll(list);
    }

    synchronized private List<Attackinfo> getAttackInfoList() {
        return this.attackinfoList;
    }

    synchronized private List<Loudongscaninfo> getScanInfoList() {
        return this.scaninfoList;
    }

    /**
     * 添加一个新扫描信息进入轮询服务
     */
    public void addScanIntoList(Loudongscaninfo info) {
        getScanInfoList().add(info);
    }

    /**
     * 添加一个打击信息进入轮询服务
     */
    public void addAttackIntoList(Attackinfo info) {
        getAttackInfoList().add(info);
    }

    /**
     * 处理扫描信息
     */
    private void processOfScanInfoList() {
        Iterator<Loudongscaninfo> iterator = getScanInfoList().iterator();
        while (iterator.hasNext()) {
            Loudongscaninfo loudongscaninfo = iterator.next();
            //如果目标的进度已经超过规定进度 那么目标移除轮询服务
            if (loudongscaninfo.getProgress() >= FULLPROGRESS) {
                loudongscaninfo.setProgress(FULLPROGRESS);
                loudongscaninfo.setFinishtime(LocalDateTime.now());
                iterator.remove();
                continue;
            }
            Loudongtargetinfo target = targetService.getById(loudongscaninfo.getTid());
            //如果扫描的目标信息丢失了的话 那么当前目标移除轮询服务
            if (target == null) {
                iterator.remove();
                continue;
            }
            //检测是否可以添加高危漏洞
            boolean highLevelSwitch = false;
            if (StoreManager.checkUrlHighLevel(configuration, target.getUrl())) {
                highLevelSwitch = true;
            }
            //尝试给扫描 添加bug
            int newBugNums = addBugIntoInfo(loudongscaninfo, highLevelSwitch);
            //尝试更新进度
            Integer newProgress = loudongscaninfo.getProgress() + MyTool.getRandomInt(80, 140);
            newProgress = newProgress > FULLPROGRESS ? FULLPROGRESS : newProgress;
            loudongscaninfo.setProgress(newProgress);
            //如果目标的新的进度已经达到了100% 那么目标状态可以设置为“已完成”
            //同时将目标移除轮询队列
            if (loudongscaninfo.getProgress() == FULLPROGRESS) {
                loudongscaninfo.setStatus(LoudongscaninfoController.STATUS_FINISH);
                loudongscaninfo.setFinishtime(LocalDateTime.now());
                iterator.remove();
            }
            LoudongScaninfoWebsocketPacket loudongScaninfoWebsocketPacket = new LoudongScaninfoWebsocketPacket(loudongscaninfo, newBugNums);
            ResultWebSocket resultWebSocket = ResultWebSocket.success(ResultWebSocket.ACTION_SCAN, loudongScaninfoWebsocketPacket);
            WebSocket.sendMessage(JSON.toJSONString(resultWebSocket));
            // 更新扫描进度信息到数据库
            scanInfoService.updateById(loudongscaninfo);
        }
    }

    /**
     * 处理打击信息
     */
    private void processOfAttackInfoList() {
        Iterator<Attackinfo> iterator = getAttackInfoList().iterator();
        while (iterator.hasNext()) {
            Attackinfo attackinfo = iterator.next();
            //生成一次攻击事件
            Activityinfo activityinfo = null;
            if (ActivityGenerator.checkCreateAvailable()) {
                activityinfo = generateActivityInfo(attackinfo);
                if (activityinfo != null) {
                    activityinfoService.save(activityinfo);
                }
            }
            //攻击的进度 进行更新
            int progress = attackinfo.getProgress();
            progress = AttackGenerator.generateNewProgress(progress);
            //设置新的攻击进度
            attackinfo.setProgress(progress);
            attackinfoService.updateById(attackinfo);
            //发送消息
            AttackWebSocketInfo attackWebSocketInfo = new AttackWebSocketInfo(attackinfo, activityinfo);
            ResultWebSocket packet = ResultWebSocket.success(ResultWebSocket.ACTION_ATTACK, attackWebSocketInfo);
            WebSocket.sendMessage(JSON.toJSONString(packet));
        }
    }

    private Activityinfo generateActivityInfo(Attackinfo attackinfo) {
        Methodinfo methodinfo = methodinfoService.getById(attackinfo.getMid());
        if (methodinfo == null) {
            //无法找到对应的方法
            return null;
        }
        Activityinfo activityinfo = new Activityinfo();
        activityinfo.setPort(ActivityGenerator.generatePort());
        activityinfo.setReason(ActivityGenerator.generateReason(methodinfo.getType()));
        activityinfo.setThread(ActivityGenerator.generateThread());
        activityinfo.setSrcip(ActivityGenerator.generateIpAddr());
        activityinfo.setDstip(attackinfo.getDstip());
        activityinfo.setAid(attackinfo.getId());
        return activityinfo;
    }

    /**
     * 添加bug信息到扫描信息中
     */
    private int addBugIntoInfo(Loudongscaninfo info, boolean highLevelBugSwitch) {
        int bugNums = 0;
        if (info.getProgress() >= FAST_INCREASE_PROGRESS) {
            if (MyTool.getRandomInt(0, 100) < 10) {
                bugNums = MyTool.getRandomInt(0, 2);
            }
        } else {
            bugNums = MyTool.getRandomInt(0, 2);
        }
        Integer targetId = info.getTid();
        Loudongtargetinfo target = targetService.getById(targetId);
        if (target == null) {
            return 0;
        }
        Loudongbuginfo[] bugs = new Loudongbuginfo[bugNums];
        for (int i = 0; i < bugs.length; i++) {
            bugs[i] = new Loudongbuginfo();
            Loudongbuginfo bug = bugs[i];
            int typeIndex = MyTool.getRandomInt(0, bugNames.length);
            //从预置库里挑选一个bug
            getBugGenerator().pickBug(highLevelBugSwitch);
            bug.setName(getBugGenerator().getBugNameFromPicker(typeIndex));
            bug.setStatus(LoudongbuginfoController.STATUS_OPEN);
            bug.setTid(targetId);
            bug.setSid(info.getId());
            bug.setParam(getBugGenerator().getParamFromPicker());
            bug.setLevel(getBugGenerator().getLevelFromPicker());
            bug.setUrl(target.getUrl());
            bug.setType(typeIndex);
            bug.setCreatedtime(LocalDateTime.now());
            bugService.save(bug);
        }
        return bugNums;
    }

    /**
     * 攻击方法
     */
    private enum MethodType {
        worm(1), hardware(2), disk(3), system(4), refuse(5), backdoor(6);

        int value;

        MethodType(int val) {
            this.value = val;
        }

        void setValue(int val) {
            this.value = val;
        }

        MethodType valueOf() {
            switch (value) {
                case 1:
                    return worm;
                case 2:
                    return hardware;
                case 3:
                    return disk;
                case 4:
                    return system;
                case 5:
                    return refuse;
                case 6:
                    return backdoor;
                default:
                    return null;
            }
        }
    }

    private static class AttackGenerator {
        private static final int PROGRESS_MAX = 9;
        private static final int PROGRESS_MIN = 4;

        public static int generateNewProgress(int oldProgress) {
            return oldProgress + MyTool.getRandomInt(PROGRESS_MIN, PROGRESS_MAX);
        }
    }

    private static class ActivityGenerator {
        private static final String wormReasons[] = {
                "smb蠕虫",
                "ftp蠕虫",
                "snmp蠕虫",
                "telnet蠕虫"
        };

        private static final String diskReasons[] = {
                "填零复写",
                "高频擦写",
                "数据破坏",
                "数据填充"
        };
        private static final String hardwareReasons[] = {
                "硬件破坏",
                "驱动破坏",
                "bios破坏",
                "硬盘摧毁"
        };
        private static final String systemReasons[] = {
                "windows_7系统破坏病毒",
                "windows_10系统破坏病毒",
                "centos7系统破坏病毒",
                "ubuntu12系统破坏病毒"
        };
        private static final String digitalReasons[] = {
                "数字大炮"
        };
        private static final String backdoorReasons[] = {
                "rat后门"
        };

        public static Integer generatePort() {
            return MyTool.getRandomInt(1000, 60000);
        }

        public static String generateReason(int type) {
            MethodType methodType = MethodType.worm;
            methodType.setValue(type);
            methodType = methodType.valueOf();
            if (methodType != null) {
                switch (methodType) {
                    case worm:
                        return wormReasons[MyTool.getRandomInt(0, wormReasons.length)];
                    case hardware:
                        return hardwareReasons[MyTool.getRandomInt(0, hardwareReasons.length)];
                    case disk:
                        return diskReasons[MyTool.getRandomInt(0, digitalReasons.length)];
                    case refuse:
                        return digitalReasons[MyTool.getRandomInt(0, digitalReasons.length)];
                    case system:
                        return systemReasons[MyTool.getRandomInt(0, systemReasons.length)];
                    case backdoor:
                        return backdoorReasons[MyTool.getRandomInt(0, backdoorReasons.length)];
                    default:
                        return null;
                }
            }
            return null;
        }

        public static Integer generateThread() {
            return MyTool.getRandomInt(1, 7);
        }

        public static String generateIpAddr() {
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < 3; i++) {
                sb.append(randomIpValue());
                sb.append(".");
            }
            sb.append(randomIpValue());
            return sb.toString();
        }

        private static String randomIpValue() {
            int randomInt = MyTool.getRandomInt(10, 255);
            return String.valueOf(randomInt);
        }

        /**
         * @return true 代表可以进行创建攻击事件
         */
        public static boolean checkCreateAvailable() {
            int randomInt = MyTool.getRandomInt(0, 100);
            if (randomInt < 10) {
                return true;
            }
            return false;
        }
    }

    private class BugGenerator {

        /**
         * 预设BUG列表[从数据库中读取]
         */
        private Presetbuginfo[] arr;
        private int cursor;
        private boolean highLevelBugSwitch;

        private BugGenerator() {
            System.out.println("从数据库读取预置漏洞库");
            List<Presetbuginfo> list = presetbugService.list();
            arr = new Presetbuginfo[list.size()];
            list.toArray(arr);
        }

        private void pickBug(boolean highLevelBugSwitch) {
            this.highLevelBugSwitch = highLevelBugSwitch;
            cursor = MyTool.getRandomInt(0, arr.length);
        }


        /**
         * 随机生成几个Bug等级
         */
        private Integer getLevelFromPicker() {
            if (arr.length == 0) {
                return 0;
            }
            Presetbuginfo presetbuginfo = arr[cursor];
            Integer level = presetbuginfo.getLevel();
            if (!highLevelBugSwitch && level > 3) {
                level = 3;
            }
            return level;
        }

        /**
         * 生成bug的参数
         */
        private String getParamFromPicker() {
            if (arr.length == 0) {
                return "无数据";
            }
            Presetbuginfo presetbuginfo = arr[cursor];
            return presetbuginfo.getParam();
        }

        /**
         * 生成bug的名称
         */
        private String getBugNameFromPicker(Integer index) {
            if (arr.length == 0) {
                return "无数据";
            }
            Presetbuginfo presetbuginfo = arr[cursor];
            return presetbuginfo.getName();
        }

    }

    @Data
    private class LoudongScaninfoWebsocketPacket implements Serializable {
        Loudongscaninfo info;
        Integer bugNums;

        public LoudongScaninfoWebsocketPacket(Loudongscaninfo loudongscaninfo, int newBugNums) {
            this.info = loudongscaninfo;
            this.bugNums = newBugNums;
        }
    }

    @Data
    private class AttackWebSocketInfo implements Serializable {
        Attackinfo attackinfo;
        Activityinfo activityinfo;

        public AttackWebSocketInfo(Attackinfo attackinfo, Activityinfo activityinfo) {
            this.attackinfo = attackinfo;
            this.activityinfo = activityinfo;
        }
    }
}
